import numpy as np
import pymc3 as pm
import theano.tensor as tt



def mcmc(N,I,rt):
    """根据输入数据,实现参数估计

    Args:
        N (int): 被试数量
        I (int): 题目数量
        rt (array)): 作答时间,shape=(N,I)
    """
    def l_s(z_square,I):
        """计算l_s"""
        c1 = (z_square.sum(axis=1)/I)**(1/3)
        c2 = (1-2/(9*I))
        c3 = (2/(9*I))**0.5
        return (c1-c2)/c3
    def l_0(z_square,epsilon):
        """计算l_0"""
        c1=z_square
        c2 = tt.log(2*np.pi*(epsilon**2))
        return (c1+c2).sum(axis=1)
    def l_z(z_square,I):
        """计算l_z"""
        c1 = z_square.sum(axis=1)
        c2 = I
        c3 = (2*I)**0.5
        return (c1-c2)/c3
    with pm.Model() as model:
        # 给定先验
        zeta_pre = pm.Normal("zeta",0,1,shape=(N,1))
        lambda_mean = pm.Normal("lambda_mean",4,2)
        lambda_scale = pm.InverseGamma("lambda_scale",1,1)
        lambda_pre = pm.Normal("lambda",lambda_mean,lambda_scale,shape=(1,I))
        phi_mean = pm.Normal("phi_mean",1,2)
        phi_scale = pm.InverseGamma("phi_scale",1,1)
        phi_pre = pm.Normal("phi",phi_mean,phi_scale,shape=(1,I))
        epsilon_pre = pm.InverseGamma("epsilon",1,1,shape=(1,I))
        mu_ip = phi_pre*(lambda_pre-zeta_pre)
        # 计算和储存拟合指标
        z_squre = ((tt.log(rt)-mu_ip)/epsilon_pre)**2 #  ((np.log(rt)-phi_i*(lambda_i-zeta_p))/epsilon)**2
        l_0 = l_0(z_squre,epsilon_pre)
        l_z = l_z(z_squre,I)
        l_s = l_s(z_squre,I)
        pm.Deterministic("z_squre",z_squre)
        pm.Deterministic("l_0",l_0)
        pm.Deterministic("l_z",l_z)
        pm.Deterministic("l_s",l_s)
        pm.Lognormal("T",mu_ip,epsilon_pre,observed=rt)
        # 进行采样
        trace = pm.sample(5000,tune=5000,chains=2,
                            return_inferencedata=True,
                            cores=4)
    return trace